# -*- coding: utf-8 -*-
from typing import Optional, Any, List, Dict

from loguru import logger
from pydantic import Field, BaseModel
from sentence_transformers import SentenceTransformer


class Embeddings(BaseModel):
    """增强的嵌入模型类"""
    model_name: str
    client: Optional[Any] = None
    cache_folder: Optional[str] = None
    model_kwargs: Dict[str, Any] = Field(default_factory=dict)
    encode_kwargs: Dict[str, Any] = Field(default_factory=dict)
    multi_process: bool = False
    show_progress: bool = False
    embeddings_cache: Optional[Any] = None

    def __init__(self, **kwargs: Any):
        """初始化嵌入模型"""
        super().__init__(**kwargs)
        try:
            self.client = SentenceTransformer(
                self.model_name,
                cache_folder=self.cache_folder,
                **self.model_kwargs
            )
        except Exception as e:
            logger.error(f"Failed to initialize SentenceTransformer: {e}")
            raise

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        if not texts:
            return []

        try:
            embeddings = self.client.encode(
                texts,
                show_progress_bar=self.show_progress,
                batch_size=32,
                **self.encode_kwargs
            )
            return embeddings.tolist()
        except Exception as e:
            logger.error(f"Embedding failed: {e}")
            raise

    def embed_query(self, text: str) -> List[float]:
        return self.embed_documents([text])[0]
